package smartai.examples.speech.tts;


import ai.djl.modality.audio.Audio;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.speech.tts.config.TtsModelConfig;
import cn.smartjavaai.speech.tts.entity.SherpaTtsParams;
import cn.smartjavaai.speech.tts.enums.TtsModelEnum;
import cn.smartjavaai.speech.tts.factory.TtsModelFactory;
import cn.smartjavaai.speech.tts.model.SherpaTtsModel;
import cn.smartjavaai.speech.tts.model.TtsModel;
import cn.smartjavaai.speech.utils.AudioUtils;
import com.k2fsa.sherpa.onnx.GeneratedAudio;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * 语音合成demo（TTS）
 * sherpa-onnx模型及依赖库下载链接:
 * 1、（推荐）依赖库官网下载：https://github.com/k2-fsa/sherpa-onnx/releases
 * 2、（推荐）TTS模型官网下载：https://github.com/k2-fsa/sherpa-onnx/releases/tag/tts-models
 * 3、百度网盘下载：https://pan.baidu.com/s/186REUf7p1z0HH9AZNnwCUg?pwd=1234 提取码: 1234
 * @author dwj
 * @date 2025/10/22
 */
@Slf4j
public class TtsDemo {

    public static String enText = "Today as always, men fall into two groups: slaves and free men. Whoever does not have"
            + " two-thirds of his day for himself, is a slave, whatever he may be: a statesman, a"
            + " businessman, an official, or a scholar.";

    public static String zhText = "有问题，请拨打110或者手机18601239876。我们的价值观是真诚热爱！";

    String znEnText =
            "中英文语音合成测试。This is generated by next generation Kaldi using Kokoro without Misaki."
                    + " 你觉得中英文说的如何呢？";

    //设备类型
    public static DeviceEnum device = DeviceEnum.CPU;


    /**
     * 获取Vits中文模型
     * @return
     */
    public TtsModel getVitsZhModel(){
        TtsModelConfig modelConfig = new TtsModelConfig();
        modelConfig.setModelEnum(TtsModelEnum.SHERPA_VITS);
        modelConfig.setModelPath("/Users/wenjie/Documents/develop/model/speech/tts/vits-zh-hf-keqing");
        modelConfig.setModelName("keqing.onnx");
        modelConfig.setLibPath(Paths.get("/Users/wenjie/smartjavaai_cache/sherpa-onnx-v1.12.14-osx-arm64-jni/lib"));
        modelConfig.setDevice(device);
        modelConfig.putCustomParam("debug", false);
        modelConfig.putCustomParam("numThreads", 1);
        return TtsModelFactory.getInstance().getModel(modelConfig);
    }

    /**
     * 获取Matcha中文模型
     * @return
     */
    public TtsModel getMatchaZhModel(){
        TtsModelConfig modelConfig = new TtsModelConfig();
        modelConfig.setModelEnum(TtsModelEnum.SHERPA_MATCHA);
        modelConfig.setModelPath("/Users/wenjie/Documents/develop/model/speech/tts/matcha-icefall-zh-baker");
        modelConfig.setModelName("model-steps-3.onnx");
        modelConfig.putCustomParam("vocoder", "/Users/wenjie/Documents/develop/model/speech/tts/matcha-icefall-zh-baker/vocos-22khz-univ.onnx");
        modelConfig.setLibPath(Paths.get("/Users/wenjie/smartjavaai_cache/sherpa-onnx-v1.12.14-osx-arm64-jni/lib"));
        modelConfig.setDevice(device);
        modelConfig.putCustomParam("debug", false);
        modelConfig.putCustomParam("numThreads", 1);
        return TtsModelFactory.getInstance().getModel(modelConfig);
    }

    /**
     * 获取Kokoro多语言模型
     * @return
     */
    public TtsModel getKokoroZhEnModel(){
        TtsModelConfig modelConfig = new TtsModelConfig();
        modelConfig.setModelEnum(TtsModelEnum.SHERPA_KOKORO);
        modelConfig.setModelPath("/Users/wenjie/Documents/develop/model/speech/tts/kokoro-multi-lang-v1_0");
        modelConfig.setModelName("model.onnx");
        modelConfig.setLibPath(Paths.get("/Users/wenjie/smartjavaai_cache/sherpa-onnx-v1.12.14-osx-arm64-jni/lib"));
        modelConfig.setDevice(device);
        modelConfig.putCustomParam("debug", false);
        modelConfig.putCustomParam("numThreads", 2);
        return TtsModelFactory.getInstance().getModel(modelConfig);
    }

    /**
     * 获取Matcha模型（英文）
     * @return
     */
    public TtsModel getMatchaEnModel(){
        TtsModelConfig modelConfig = new TtsModelConfig();
        modelConfig.setModelEnum(TtsModelEnum.SHERPA_MATCHA);
        modelConfig.setModelPath("/Users/wenjie/Documents/develop/model/speech/tts/matcha-icefall-en_US-ljspeech");
        modelConfig.setModelName("model-steps-3.onnx");
        modelConfig.setLibPath(Paths.get("/Users/wenjie/smartjavaai_cache/sherpa-onnx-v1.12.14-osx-arm64-jni/lib"));
        modelConfig.putCustomParam("vocoder", "/Users/wenjie/Documents/develop/model/speech/tts/matcha-icefall-zh-baker/vocos-22khz-univ.onnx");
        modelConfig.setDevice(device);
        modelConfig.putCustomParam("debug", false);
        modelConfig.putCustomParam("numThreads", 2);
        return TtsModelFactory.getInstance().getModel(modelConfig);
    }

    /**
     * 语音合成Vits（中文）
     */
    @Test
    public void testVitsZhTts() throws IOException {
        TtsModel ttsModel = getVitsZhModel();
        SherpaTtsParams params = new SherpaTtsParams();
        //语速
        params.setSpeed(1f);
        //说话人ID
        params.setSpeakerId(100);
        //生成音频到本地路径
        ttsModel.generate(zhText, params, "/Users/wenjie/Downloads/tts-vits-zh.wav");
    }

    /**
     * 语音合成Matcha（中文）
     */
    @Test
    public void testMatchaZhTts() throws IOException {
        TtsModel ttsModel = getMatchaZhModel();
        SherpaTtsParams params = new SherpaTtsParams();
        //语速
        params.setSpeed(1f);
        //说话人ID
        params.setSpeakerId(0);
//        ttsModel.generate(zhText, params,"/Users/wenjie/Downloads/tts-matcha-zh.wav");
        R<Audio> result = ttsModel.generate(znEnText, params);
        if (result.isSuccess()){
            Audio audio = result.getData();
            AudioUtils.saveToWav(audio,"/Users/wenjie/Downloads/tts-matcha-zh.wav");
        }else{
            System.out.println(result.getMessage());
        }


    }

    /**
     * 语音合成Kokoro（中英）
     */
    @Test
    public void testKokoroZhEnTts() throws IOException {
        TtsModel ttsModel = getKokoroZhEnModel();
        SherpaTtsParams params = new SherpaTtsParams();
        //语速
        params.setSpeed(1f);
        //说话人ID
        params.setSpeakerId(3);
//        ttsModel.generate(znEnText, params, "/Users/wenjie/Downloads/tts-kokoro-zh-en.wav");
        R<Audio> result = ttsModel.generate(znEnText, params);
        if (result.isSuccess()){
            Audio audio = result.getData();
            AudioUtils.saveToWav(audio,"/Users/wenjie/Downloads/tts-kokoro-zh-en.wav");
        }else{
            System.out.println(result.getMessage());
        }
    }

    /**
     * 语音合成Matcha（英文）
     */
    @Test
    public void testMatchaEnTts() throws IOException {
        TtsModel ttsModel = getMatchaEnModel();
        SherpaTtsParams params = new SherpaTtsParams();
        //语速
        params.setSpeed(1f);
        //说话人ID
        params.setSpeakerId(0);
//        ttsModel.generate(enText, params, "/Users/wenjie/Downloads/tts-kitten-en.wav");
        R<Audio> result = ttsModel.generate(enText, params);
        if (result.isSuccess()){
            Audio audio = result.getData();
            AudioUtils.saveToWav(audio,"/Users/wenjie/Downloads/tts-kitten-en.wav");
        }else{
            System.out.println(result.getMessage());
        }
    }






}
